"""
@Time: 2021/3/10 下午 9:58
@Author: jinzhuan
@File: BBNb.py
@Desc: 
"""
import sys

sys.path.append('/data/zhuoran/code/cognlp')
sys.path.append('/data/zhuoran/cognlp')

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import RandomSampler
from cognlp import *

torch.cuda.set_device(0)
device = torch.device('cuda')

from cognlp.io.loader.et.BBNb import BBNbEtLoader
from cognlp.io.processor.et.BBNb import BBNbEtProcessor
from cognlp.models.et.bert_et_binary import Bert4EtBinary

loader = BBNbEtLoader()
train_data, test_data = loader.load_all('../../../cognlp/data/et/BBN/data')

processor = BBNbEtProcessor(label_list=loader.get_labels(), path='../../../cognlp/data/et/BBN/data',
                            bert_model='bert-base-cased')
vocabulary = Vocabulary.load('../../../cognlp/data/et/BBN/data/vocabulary.txt')

# train_datable = processor.process(train_data)
# train_dataset = DataTableSet(train_datable)
# train_sampler = RandomSampler(train_dataset)

test_datable = processor.process(test_data)
test_dataset = DataTableSet(test_datable)
test_sampler = RandomSampler(test_dataset)

model = Bert4EtBinary(len(vocabulary), bert_model='bert-base-cased', embedding_size=768)
# metric = AccuracyMetric()
metric = ClassifyFPreRecMetric(f_type='macro')
# metric = ClassifyFPreRecMetric(f_type='micro')

tester = Tester(model, model_path='../../../cognlp/data/et/BBN/model/BBNb/2021-03-10-22:20:20/checkpoint-2614/model.pt',
                batch_size=20,
                sampler=test_sampler,
                drop_last=False, num_workers=5, print_every=1000,
                dev_data=test_dataset, metrics=metric, metric_key=None, use_tqdm=True, device=device,
                callbacks=None, check_code_level=0, device_ids=[0])
tester.test()
# loss = nn.CrossEntropyLoss()
# optimizer = optim.Adam(model.parameters(), lr=0.00001)
# trainer = Trainer(model,
#                   train_dataset,
#                   dev_data=test_dataset,
#                   n_epochs=20,
#                   batch_size=100,
#                   loss=loss,
#                   optimizer=optimizer,
#                   scheduler=None,
#                   metrics=metric,
#                   train_sampler=train_sampler,
#                   dev_sampler=test_sampler,
#                   drop_last=False,
#                   gradient_accumulation_steps=1,
#                   num_workers=5,
#                   save_path="../../../cognlp/data/et/BBN/model",
#                   save_file=None,
#                   print_every=None,
#                   scheduler_steps=None,
#                   validate_steps=100,
#                   save_steps=None,
#                   grad_norm=None,
#                   use_tqdm=True,
#                   device=device,
#                   device_ids=[0, 1, 2, 3, 4, 5],
#                   callbacks=None,
#                   metric_key=None,
#                   writer_path='../../../cognlp/data/et/BBN/tensorboard',
#                   fp16=False,
#                   fp16_opt_level='O1',
#                   checkpoint_path=None,
#                   task='BBNb',
#                   logger_path='../../../cognlp/data/et/BBN/logger')
#
# trainer.train()
